#!/usr/bin/env python3
import os
import tarfile
import shutil
import argparse
import paramiko
from datetime import datetime

def create_backup(source_dir, backup_dir):
        """
        将 source_dir 打包为.tar.gz文件，存放在 backup_dir 中，返回打包后的文件路径
        """

        # 创建备份文件名（含时间戳）

        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        backup_filename = f"backup_{timestamp}.tar.gz"
        backup_path = os.path.join(backup_dir, backup_filename)

        print(f"[+] 开始打包目录 {source_dir} 到 {backup_path}")

        # 打包压缩

        with tarfile.open(backup_path, "w:gz") as tar:
                tar.add(source_dir, arcname=os.path.basename(source_dir))
        print("[+] 打包完成")
        return backup_path


def upload_to_remote(local_file, remote_host, remote_user, remote_pass, remote_path):
        """
        使用SSH/SFTP 将本地文件上传到远程服务器
        """
        print(f"[+] 正在连接远程服务器 {remote_host}")
        transport = paramiko.Transport((remote_host, 10048))
        transport.connect(username=remote_user, password=remote_pass)
        sftp = paramiko.SFTPClient.from_transport(transport)

        print(f"[+] 正在上传 {local_file} 到 {remote_host}:{remote_path}")
        remote_file = os.path.join(remote_path, os.path.basename(local_file))
        sftp.put(local_file, remote_file)

        print("[+] 上传完成")

        #关闭连接

        sftp.close
        transport.close

        return remote_file

def extract_on_remote(remote_host, remote_user, remote_pass, remote_file):
        """
        在远程服务器上执行解压命令
        """

        print(f"[+] 正在登陆远程服务器执行解压命令...")
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(remote_host, port=10048, username=remote_user, password=remote_pass)

        #获取远程文件目录
        remote_dir = os.path.dirname(remote_file)
        cmd = f"cd {remote_dir} && tar -xzf {remote_file}"
        stdin, stdout, stderr = ssh.exec_command(cmd)


        #输出执行结果
        for line in stdout.readlines():
                print("输出:", line.strip())
        for line in stderr.readlines():
                print("错误:", line.strip())

        print("[+] 解压完成")
        ssh.close()

def main():
        parser = argparse.ArgumentParser(description="自动备份目录并上传到远程服务器")
        parser.add_argument("--source", required=True, help="要备份的本地目录路径")
        parser.add_argument("--remote-host", required=True, help="远程服务器 IP 地址")
        parser.add_argument("--remote-user", required=True, help="远程服务器用户名")
        parser.add_argument("--remote-pass", required=True, help="远程服务器密码")
        parser.add_argument("--remote-path", required=True, help="远程服务器保存路径")
        parser.add_argument("--no-extract", action="store_true", help="不自动解压上传的备份文件")
        args = parser.parse_args()

        # 配置本地临时备份目录
        local_backup_dir = "/tmp/backup"

        if not os.path.exists(local_backup_dir):
                os.makedirs(local_backup_dir)

        try:
                #1. 打包本地目录
                backup_file = create_backup(args.source, local_backup_dir)

                #2. 上传到远程服务器
                remote_file = upload_to_remote(
                                backup_file,
                                args.remote_host,
                                args.remote_user,
                                args.remote_pass,
                                args.remote_path
                                )

                # 3. 是否在远程服务器解压
                if not args.no_extract:
                        extract_on_remote(
                                        args.remote_host,
                                        args.remote_user,
                                        args.remote_pass,
                                        remote_file
                                        )
                print("[+] 备份任务上传完成")
        except Exception as e:
                print("[-] 发生异常:", str(e))
        finally:
                #清理本地临时文件
                if os.path.exists(backup_file):
                        os.remove(backup_file)
                        print("[+] 已清理本地临时备份文件")

if __name__ == "__main__":
        main()